"""
@Time: 2021/1/18 下午 7:31
@Author: jinzhuan
@File: fn_toolkit.py
@Desc: 
"""
import threading
import cognlp.io.processor.fn.framenet as processor
import torch
from cognlp import *
from ..base_toolkit import BaseToolkit


class FnToolkit(BaseToolkit):

    def __init__(self, bert_model=None, model_path=None, vocabulary_path=None, device=None,
                 device_ids=None, max_seq_length=None):
        super().__init__(bert_model, model_path, vocabulary_path, device, device_ids, max_seq_length)
        self.model = Bert4Frame(len(self.vocabulary))
        self.load_model()

    _instance_lock = threading.Lock()

    def __new__(cls, *args, **kwargs):
        if not hasattr(FnToolkit, "_instance"):
            with FnToolkit._instance_lock:
                if not hasattr(FnToolkit, "_instance"):
                    FnToolkit._instance = object.__new__(cls)
        return FnToolkit._instance

    def run(self, words):
        self.model.eval()
        frames = []

        input_ids, attention_mask, head_indexes, frame_id, element_id, label_mask = \
            processor.process(words, [], [], self.tokenizer, self.vocabulary, None, self.max_seq_length)
        with torch.no_grad():
            prediction, valid_len = self.model.predict(
                [[input_ids], [attention_mask], [head_indexes], [frame_id], [element_id], [label_mask]])
        if len(prediction) == 0:
            return []
        prediction = prediction[0]
        valid_len = valid_len[0]

        for i in range(valid_len.item()):
            if prediction[i].item() != self.vocabulary.to_index("<unk>"):
                frame = {"word": words[i],
                         "position": i,
                         "frame": self.vocabulary.to_word(prediction[i].item())}
                frames.append(frame)
        return frames

"""
Input:
The true voodoo-worshipper attempts nothing of importance without certain sacrifices which are intended to propitiate his unclean gods.
Output:
[
    {
        "word": "attempts",
        "position": 3,
        "frame": "Attempt"
    },
    {
        "word": "importance",
        "position": 6,
        "frame": "Importance"
    },
    {
        "word": "sacrifices",
        "position": 9,
        "frame": "Rite"
    },
    {
        "word": "intended",
        "position": 12,
        "frame": "Purpose"
    }
]
"""